from pspnet_voc import PSPNet
from torch import nn
from PIL import Image
from torch.autograd import Variable
import torch.nn.functional as F
import numpy as np
import colorsys
import torch
import copy
import os
import cv2
import numpy
import random


def create_visual_anno(anno):
    """"""
    print(np.max(anno))
    # assert np.max(anno) <= 21, "only 21 classes are supported, add new color in label2color_dict"
    # label2color_dict = {
    #     0: [random.randint(0,100), random.randint(0,100), random.randint(0,100)],
    #     1: [random.randint(0,100), random.randint(0,100), random.randint(0,100)],
    #     2: [random.randint(0,100), random.randint(0,100), random.randint(0,100)],
    #     3: [random.randint(0,100), random.randint(0,100), random.randint(0,100)],
    #     4: [random.randint(0,100), random.randint(0,100), random.randint(0,100)],
    #     5: [random.randint(0,100), random.randint(0,100), random.randint(0,100)],
    #     6: [random.randint(0,100), random.randint(0,100), random.randint(0,100)],
    #     7: [random.randint(0,100), random.randint(100,255), random.randint(0,255)],
    #     8: [random.randint(0,100), random.randint(100,255), random.randint(0,255)],
    #     9: [random.randint(0,100), random.randint(100,255), random.randint(0,255)],
    #     10: [random.randint(0,100), random.randint(100,255), random.randint(0,255)],
    #     11: [random.randint(0,100), random.randint(100,255), random.randint(0,255)],
    #     12: [random.randint(0,100), random.randint(100,255), random.randint(0,255)],
    #     13: [random.randint(0,255), random.randint(100,255), random.randint(0,255)],
    #     14: [random.randint(0,255), random.randint(100,255), random.randint(0,255)],
    #     15: [random.randint(0,255), random.randint(100,255), random.randint(0,255)],
    #     16: [random.randint(0, 255), random.randint(100, 255), random.randint(0, 255)],
    #     17: [random.randint(0, 255), random.randint(100, 255), random.randint(0, 255)],
    #     18: [random.randint(0, 255), random.randint(100, 255), random.randint(0, 255)],
    #     19: [random.randint(0, 100), random.randint(100, 255), random.randint(0, 255)],
    #     20: [random.randint(0, 100), random.randint(100, 255), random.randint(0, 255)],
    #     21: [random.randint(0, 100), random.randint(100, 255), random.randint(0, 255)]
    # }
    label2color_dict=[]
    label2color_dict.append([0, 0, 0])
    for i in range(1, 256):
        label2color_dict.append([random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)])

    # visualize
    visual_anno = np.zeros((anno.shape[0], anno.shape[1], 3), dtype=np.uint8)
    for i in range(visual_anno.shape[0]):  # i for h
        for j in range(visual_anno.shape[1]):
            cls = anno[i, j]
            color = label2color_dict[anno[i, j]]
            visual_anno[i, j, 0] = color[0]
            visual_anno[i, j, 1] = color[1]
            visual_anno[i, j, 2] = color[2]

    return visual_anno


class miou_Pspnet(PSPNet):
    def detect_image(self, image):
        # model_dict = self.net.state_dict()
        # pretrained_dict = torch.load("Epoch29-Total_Loss0.5295-Val_Loss0.4825.pth")
        # pretrained_dict = {k: v for k, v in pretrained_dict.items() if np.shape(model_dict[k]) == np.shape(v)}
        # model_dict.update(pretrained_dict)
        # self.net.load_state_dict(model_dict)
        orininal_h = np.array(image).shape[0]
        orininal_w = np.array(image).shape[1]

        image, nw, nh = self.letterbox_image(image, (self.model_image_size[1], self.model_image_size[0]))
        images = [np.array(image) / 255]
        images = np.transpose(images, (0, 3, 1, 2))

        with torch.no_grad():
            images = Variable(torch.from_numpy(images).type(torch.FloatTensor))
            if self.cuda:
                images = images.cuda()
            pr = self.net(images)[0]
            pr = F.softmax(pr.permute(1, 2, 0), dim=-1).cpu().numpy().argmax(axis=-1)

        pr = pr[int((self.model_image_size[0] - nh) // 2):int((self.model_image_size[0] - nh) // 2 + nh),
             int((self.model_image_size[1] - nw) // 2):int((self.model_image_size[1] - nw) // 2 + nw)]

        image = Image.fromarray(np.uint8(pr)).resize((orininal_w, orininal_h), Image.NEAREST)

        return image


pspnet = miou_Pspnet()

image_ids = open(r"VOCdevkit\VOC2007\ImageSets\Segmentation\val.txt", 'r').read().splitlines()

if not os.path.exists("./miou_pr_dir"):
    os.makedirs("./miou_pr_dir")

for image_id in image_ids:
    image_path = "./VOCdevkit/VOC2007/JPEGImages/" + image_id + ".jpg"
    print(image_path)
    image = Image.open(image_path)
    image = pspnet.detect_image(image)
    save_path = "./miou_pr_dir/" + image_id + ".png"
    image.save(save_path)
    print(image_id, " done!")
    img_np = numpy.asarray(image)
    img_np_cls = create_visual_anno(img_np)
    # img_np = cv2.imread(save_path, cv2.IMREAD_GRAYSCALE)
    cv2.imshow("img_np", img_np_cls)
    # img_cv = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
    # size_ori = img_cv.shape
    # img_np = img_np * 255
    # imgShow = np.empty((size_ori[0], size_ori[1], 3), np.uint8)
    # imgShow[:, :, 0] = img_cv
    # imgShow[:, :, 1] = img_cv
    # imgShow[:, :, 2] = np.maximum(img_cv, img_np) * 0.8 + img_cv * 0.2
    img_cv = cv2.imread(image_path, cv2.IMREAD_COLOR)
    imgShow = cv2.addWeighted(img_cv, 0.6, img_np_cls, 0.4, 0)
    cv2.imshow("result", imgShow)
    cv2.waitKey(1000)



